#include <vector>
#include <cmath>
#include <cfloat>
#include <iostream>

using namespace std;

template <typename T, 
    int H, 
    int N, 
    int D,
    int BR,
    int BC
>
void flash_attn_v1(const std::vector<T>& Q,
    const std::vector<T>& K,
    const std::vector<T>& V,
    std::vector<T>& out
) {
    constexpr int ND = N * D;

    const T* k_ptr = (T*)K.data();
    const T* v_ptr = (T*)V.data();
    const T* q_ptr = (T*)Q.data();
    T* o_ptr = (T*)out.data();

    vector<T> max(N, -FLT_MAX);
    vector<T> sum(N, 0.0f);

    for (int i = 0; i < N; i += BC) {
        // load K, V
        k_ptr = (T*)K.data() + i * D;
        v_ptr = (T*)V.data() + i * D;
        for (int j = 0; j < N; j += BR) {
            for (int r = 0; (r < BR) && (r + j < N); ++r) {
                // load Q，O
                q_ptr = (T*)Q.data() + (j + r) * D;
                o_ptr = (T*)out.data() + (j + r) * D;

                // 计算S = QKT -> [BR]
                // 同时计算每行的最大值
                T br_max = -FLT_MAX;
                vector<T> S(BC, 0.0f);
                for (int c = 0; (c < BC) && (c + i < N); ++c) {
                    // calc 
                    T value = 0.0;
                    for (int d = 0; d < D; ++d) {
                        value += q_ptr[d] * k_ptr[c * D + d];
                    }
                    S[c] = value / sqrtf(D);
                    // get row max
                    br_max = br_max > value ? br_max : value;
                }
                br_max = br_max > max[j + r] ? br_max : max[j + r];

                // softmax: 求每一行的和
                float br_sum = sum[j + r] * expf(max[j + r] - br_max);
                for (int c = 0; (c < BC) && (c + i < N); ++c) {
                    br_sum += expf(S[c] - br_max);
                }
                
                auto& P = S;
                for (int c = 0; (c < BC) && (c + i < N); ++c) {
                    P[c] = expf(S[c] - br_max) / br_sum;
                }
                // P * V + update(o)
                // [BR, BC] * [BC, D]
                float ratio = sum[j + r] / br_sum * expf(max[r + j] - br_max);
                for (int d = 0; d < D; ++d) {
                    float o_tmp = o_ptr[d] * ratio; 
                    // float o_tmp = o_ptr[d] * expf(-max[r + j] + br_max);  // error
                    for (int bc = 0; (bc < BC) && (bc + i < N); ++bc) {
                        o_tmp += P[bc] * v_ptr[bc * D + d];
                    }
                    o_ptr[d] = o_tmp;
                }
                // update sum, max
                sum[r + j] = br_sum;
                max[r + j] = br_max;
            }
        }
    }
}